/******************************************************************************

  File Name     : alg_face_detect.cpp
  Version       : Initial Draft
  Author        : tongzhilin
  Created       : 2020/8/6
  Last Modified :
  Description   : alg_face_detect_v1, use alg: facedetection_output_norm or facedetection_output_norm_conf_nhwc
  Function List :
  History       :
  1.Date        : 2020/8/6
    Author      : tongzhilin
    Modification: Created file

******************************************************************************/
#include "alg_face_detect_v1.h"
#include "alg_utils.h"

std::vector<MpMat32f> FaceDetV1::m_stvAnchor;

FaceDetV1::FaceDetV1(const char* _pModelName) : NnieAlgBasic(_pModelName, 1)
{
    INT32 C = 0;
    INT32 H = 0;
    INT32 W = 0;

    //default box size
    float pBoxSizes3[3] = {10, 16, 24};
    float pBoxSizes4[2] = {32, 48};
    float pBoxSizes5[2] = {64, 96};
    float pBoxSizes6[3] = {128, 192, 256};
    //default box step/stride
    INT32 pBoxStep3 = 8;
    INT32 pBoxStep4 = 16;
    INT32 pBoxStep5 = 32;
    INT32 pBoxStep6 = 64;

    if(m_stvAnchor.size() != 4)
    {
        //4个尺寸的default box 对应4个chn输出
        MP_ASSERT(m_spNnModel->get_output_blob_num()/2 == 4);

        m_stvAnchor.resize(m_spNnModel->get_output_blob_num()/2);
        //create 4 anchors
        for (INT32 i = 0; i < m_spNnModel->get_output_blob_num(); i+=2)
        {
            C = m_s32vOutBlob_C[i];
            H = m_s32vOutBlob_H[i];
            W = m_s32vOutBlob_W[i];
            MpMat32f stMpMat(H, W, C);
            m_stvAnchor[i/2] = stMpMat;
        }

        //get every anchor boxes
        get_anchor_box(pBoxStep3, 3, pBoxSizes3, m_stvAnchor[0]);
        get_anchor_box(pBoxStep4, 2, pBoxSizes4, m_stvAnchor[1]);
        get_anchor_box(pBoxStep5, 2, pBoxSizes5, m_stvAnchor[2]);
        get_anchor_box(pBoxStep6, 3, pBoxSizes6, m_stvAnchor[3]);
    }
}

FaceDetV1::~FaceDetV1()
{
    
}

bool FaceDetV1::get_anchor_box(INT32 _s32Step, INT32 _s32BoxNum, float *_pf32BoxSizes, MpMat32f &_stAnchor)
{
    INT32 s32InputW = m_stvInputSvpBlob[0].unShape.stWhc.u32Width;
    INT32 s32InputH = m_stvInputSvpBlob[0].unShape.stWhc.u32Height;
    INT32 C = _stAnchor.channels();
    INT32 H = _stAnchor.rows();
    INT32 W = _stAnchor.cols();
    INT32 c = 0;
    INT32 posi = 0;

    MP_ASSERT(C == _s32BoxNum * 4 || W == _s32BoxNum * 4);

    for (int h = 0; h < H; ++h) 
    {
        for (int w = 0; w < W; ++w) 
        {
            float* pAnchor = (float *)_stAnchor.raw_ptr();
            //priorbox
            for (int s = 0; s < _s32BoxNum; s++) 
            {
                c = s*4;
                float f32CenterX = (w + 0.5f) * _s32Step;
                float f32CenterY = (h + 0.5f) * _s32Step;
                float f32BoxSize = _pf32BoxSizes[s];

                posi = ((c+0) * H + h) * W + w;
                pAnchor[posi] = (f32CenterX - f32BoxSize / 2.f) / s32InputW;
                posi = ((c+1) * H + h) * W + w;
                pAnchor[posi] = (f32CenterY - f32BoxSize / 2.f) / s32InputH;
                posi = ((c+2) * H + h) * W + w;
                pAnchor[posi] = (f32CenterX + f32BoxSize / 2.f) / s32InputW;
                posi = ((c+3) * H + h) * W + w;
                pAnchor[posi] = (f32CenterY + f32BoxSize / 2.f) / s32InputH;
            }
        }
    }
    
    return true;
}

bool FaceDetV1::ab_image_post()
{
    //detect post. output format : c h w
    auto get_box_by_anchor = [&](SVP_BLOB_S& _stLocBlob, SVP_BLOB_S& _stConfBlob, MpMat32f &_stAnchor, 
        std::vector<ALG_RECT_RESULT_S> &_svtRect) 
    {
        //the diff between _stLocBlob and _stConfBlob is chn. chn_conf * 2 == chn_loc
        //chn_conf/2, face score. chn_conf/2 + 1: nonface score
        //conf is score, loc is rect
        INT32 C = _stConfBlob.unShape.stWhc.u32Chn;
        INT32 H = _stConfBlob.unShape.stWhc.u32Height;
        INT32 W = _stConfBlob.unShape.stWhc.u32Width;
        INT32 S = _stConfBlob.u32Stride / 4;
        INT32 c_loc = 0;
        int32_t* __restrict ps32ScorePtr = (int32_t*)(void*)_stConfBlob.u64VirAddr;
        int32_t* __restrict ps32RectPtr = (int32_t*)(void*)_stLocBlob.u64VirAddr;
        float* pAnchor = (float *)_stAnchor.raw_ptr();
        INT32 s32LocPosi = 0;
        float f32LocCx = 0;
        float f32LocCy = 0;
        float f32LocW = 0;
        float f32LocH = 0;
        float f32AnchorX1 = 0;
        float f32AnchorY1 = 0;
        float f32AnchorX2 = 0;
        float f32AnchorY2 = 0;
        float f32AnchorCx = 0;  //centry x
        float f32AnchorCy = 0;  //centry y
        float f32AnchorW = 0;
        float f32AnchorH = 0;
        float f32BoxCx = 0;
        float f32BoxCy = 0;
        float f32BoxW = 0;
        float f32BoxH = 0;
        float prior_variance[4] = {0.1f, 0.1f, 0.2f, 0.2f};
        ALG_RECT_RESULT_S stRect;
        for (INT32 c = 0; c < C; c+=2)
        for (INT32 h = 0; h < H; h++)
        for (INT32 w = 0; w < W; w++) 
        {
            //calc score maybe can optimize to nnie!!
            //move to nnie by tzl 2020.09.01
            #if 1
            INT32 s32Conf0Posi = ((c+0) * H + h) * S + w;
            float f32Score0 = ps32ScorePtr[s32Conf0Posi] / 4096.f;

            INT32 s32Conf1Posi = ((c+1) * H + h) * S + w;
            float f32Score1 = ps32ScorePtr[s32Conf1Posi] / 4096.f;

            float s32ScoreMax = MP_MAX(f32Score0, f32Score1);
            f32Score0 -= s32ScoreMax;
            f32Score1 -= s32ScoreMax;
            f32Score0 = expf(f32Score0);
            f32Score1 = expf(f32Score1);
            s32ScoreMax = f32Score0 + f32Score1;
            f32Score0 = f32Score0/s32ScoreMax;
            f32Score1 = f32Score1/s32ScoreMax;
            #else
            INT32 s32Conf1Posi = ((c+1) * H + h) * S + w;
            float f32Score1 = ps32ScorePtr[s32Conf1Posi] / 4096.f;
            #endif
            if (f32Score1 > m_f32ConfThreShold)
            {
                c_loc = c/2;
                s32LocPosi = ((c_loc*4 + 0) * H + h) * S + w;
                f32LocCx = ps32RectPtr[s32LocPosi] / 4096.f;
                f32AnchorX1 = pAnchor[s32LocPosi];

                s32LocPosi = ((c_loc*4 + 1) * H + h) * S + w;
                f32LocCy = ps32RectPtr[s32LocPosi] / 4096.f;
                f32AnchorY1 = pAnchor[s32LocPosi];

                s32LocPosi = ((c_loc*4 + 2) * H + h) * S + w;
                f32LocW = ps32RectPtr[s32LocPosi] / 4096.f;
                f32AnchorX2 = pAnchor[s32LocPosi];

                s32LocPosi = ((c_loc*4 + 3) * H + h) * S + w;
                f32LocH = ps32RectPtr[s32LocPosi] / 4096.f;
                f32AnchorY2 = pAnchor[s32LocPosi];

                f32AnchorCx = (f32AnchorX2 + f32AnchorX1)/2;
                f32AnchorCy = (f32AnchorY2 + f32AnchorY1)/2;
                f32AnchorW = f32AnchorX2 - f32AnchorX1;
                f32AnchorH = f32AnchorY2 - f32AnchorY1;

                //MP_LOG_INFO("xxxxxxxxxxxxxxxxxxxx 00  %f  %f  %f  %f \n", f32AnchorX1, f32AnchorY1, f32AnchorX2, f32AnchorY2);
                //MP_LOG_INFO("xxxxxxxxxxxxxxxxxxxx 11  %f  %f  %f  %f -- score : %f\n", f32AnchorCx, f32AnchorCy, f32AnchorW, f32AnchorH, f32Score1);

                f32BoxCx = prior_variance[0] * f32LocCx * f32AnchorW + f32AnchorCx;
                f32BoxCy = prior_variance[1] * f32LocCy * f32AnchorH + f32AnchorCy;
                f32BoxW = expf(prior_variance[2] * f32LocW) * f32AnchorW;
                f32BoxH = expf(prior_variance[3] * f32LocH) * f32AnchorH;
                //MP_LOG_INFO("xxxxxxxxxxxxxxxxxxxx 22  %f  %f  %f  %f -- score : %f\n", f32BoxCx, f32BoxCy, f32BoxW, f32BoxH, f32Score1);

                stRect.m_stRect.m_f32X1 = f32BoxCx - 0.5f * f32BoxW;
                stRect.m_stRect.m_f32Y1 = f32BoxCy - 0.5f * f32BoxH;
                stRect.m_stRect.m_f32X2 = f32BoxCx + 0.5f * f32BoxW;
                stRect.m_stRect.m_f32Y2 = f32BoxCy + 0.5f * f32BoxH;

                stRect.m_stRect.m_f32X1 = MP_MAX(0, stRect.m_stRect.m_f32X1);
                stRect.m_stRect.m_f32Y1 = MP_MAX(0, stRect.m_stRect.m_f32Y1);
                stRect.m_stRect.m_f32X2 = MP_MIN(1.f, stRect.m_stRect.m_f32X2);
                stRect.m_stRect.m_f32Y2 = MP_MIN(1.f, stRect.m_stRect.m_f32Y2);

                
                //MP_LOG_INFO("xxxxxxxxxxxxxxxxxxxx 33  %f  %f  %f  %f -- score : %f\n", 
                //    stRect.m_stRect.m_f32X1, stRect.m_stRect.m_f32Y1, 
                //    stRect.m_stRect.m_f32X2, stRect.m_stRect.m_f32Y2, f32Score1);

                stRect.m_stRect.m_f32X1 = stRect.m_stRect.m_f32X1 * 1024 + 0.5;
                stRect.m_stRect.m_f32Y1 = stRect.m_stRect.m_f32Y1 * 576 + 0.5;
                stRect.m_stRect.m_f32X2 = stRect.m_stRect.m_f32X2 * 1024 + 0.5;
                stRect.m_stRect.m_f32Y2 = stRect.m_stRect.m_f32Y2 * 576 + 0.5;
                stRect.m_f32Score = f32Score1;

                //MP_LOG_INFO("xxxxxxxxxxxxxxxxxxxx 44  %f  %f  %f  %f -- score : %f\n", 
                //    stRect.m_stRect.m_f32X1, stRect.m_stRect.m_f32Y1, 
                //    stRect.m_stRect.m_f32X2, stRect.m_stRect.m_f32Y2, f32Score1);

                if ((stRect.m_stRect.m_f32X2 - stRect.m_stRect.m_f32X1) >= 1 &&
                    (stRect.m_stRect.m_f32Y2 - stRect.m_stRect.m_f32Y1) >= 1)
                {
                    _svtRect.emplace_back(ALG_RECT_RESULT_S{ stRect });
                }
            }
        }
    };

    get_box_by_anchor(m_stvOutputSvpBlob[0], m_stvOutputSvpBlob[1], m_stvAnchor[0], m_stvRectResult);
    get_box_by_anchor(m_stvOutputSvpBlob[2], m_stvOutputSvpBlob[3], m_stvAnchor[1], m_stvRectResult);
    get_box_by_anchor(m_stvOutputSvpBlob[4], m_stvOutputSvpBlob[5], m_stvAnchor[2], m_stvRectResult);
    get_box_by_anchor(m_stvOutputSvpBlob[6], m_stvOutputSvpBlob[7], m_stvAnchor[3], m_stvRectResult);

    MP_LOG_DBG("box number before nms:%d \n", m_stvRectResult.size());

    rect_nms(m_stvRectResult, m_f32NmsThreShold);

    return true;
}

bool FaceDetV1::ab_image_post_v2()
{
    //detect post. output format :
    //Loc: c h w.
    //Conf: h w c. note:hinnie output_c is h, output_h is w, output_h is c
    auto get_box_by_anchor = [&](SVP_BLOB_S& _stLocBlob, SVP_BLOB_S& _stConfBlob, MpMat32f &_stAnchor, 
        std::vector<ALG_RECT_RESULT_S> &_svtRect) 
    {
        //conf is score, loc is rect
        INT32 H = _stConfBlob.unShape.stWhc.u32Chn;
        INT32 W = _stConfBlob.unShape.stWhc.u32Height;
        INT32 C = _stConfBlob.unShape.stWhc.u32Width;
        INT32 S_c = _stConfBlob.u32Stride / 4;
        INT32 H_loc = _stLocBlob.unShape.stWhc.u32Height;
        INT32 W_loc = _stLocBlob.unShape.stWhc.u32Width;
        INT32 C_loc = _stLocBlob.unShape.stWhc.u32Chn;
        INT32 S_loc = _stLocBlob.u32Stride / 4;
        INT32 c_loc = 0;
        int32_t* __restrict ps32ScorePtr = (int32_t*)(void*)_stConfBlob.u64VirAddr;
        int32_t* __restrict ps32RectPtr = (int32_t*)(void*)_stLocBlob.u64VirAddr;
        float* pAnchor = (float *)_stAnchor.raw_ptr();
        INT32 s32LocPosi = 0;
        //INT32 s32AnchorPosi = 0;
        float f32LocCx = 0;
        float f32LocCy = 0;
        float f32LocW = 0;
        float f32LocH = 0;
        float f32AnchorX1 = 0;
        float f32AnchorY1 = 0;
        float f32AnchorX2 = 0;
        float f32AnchorY2 = 0;
        float f32AnchorCx = 0;  //centry x
        float f32AnchorCy = 0;  //centry y
        float f32AnchorW = 0;
        float f32AnchorH = 0;
        float f32BoxCx = 0;
        float f32BoxCy = 0;
        float f32BoxW = 0;
        float f32BoxH = 0;
        float prior_variance[4] = {0.1f, 0.1f, 0.2f, 0.2f};
        ALG_RECT_RESULT_S stRect;

        MP_ASSERT(H == H_loc && W == W_loc && C*2 == C_loc);
        for (INT32 h = 0; h < H; h++)
        for (INT32 w = 0; w < W; w++)
        for (INT32 c = 0; c < C; c+=2)
        {
            //calc score maybe can optimize to nnie!!
            //move to nnie by tzl 2020.09.01
            INT32 s32Conf1Posi = (h * W + w) * S_c + c + 1;
            float f32Score1 = ps32ScorePtr[s32Conf1Posi] / 4096.f;
            if (f32Score1 > m_f32ConfThreShold)
            {
                c_loc = c/2;
                s32LocPosi = ((c_loc*4 + 0) * H + h) * S_loc + w;
                //s32AnchorPosi = ((c_loc*4 + 0) * H + h) * S + w;
                f32LocCx = ps32RectPtr[s32LocPosi] / 4096.f;
                f32AnchorX1 = pAnchor[s32LocPosi];

                s32LocPosi = ((c_loc*4 + 1) * H + h) * S_loc + w;
                //s32AnchorPosi = ((c_loc*4 + 1) * H + h) * S + w;
                f32LocCy = ps32RectPtr[s32LocPosi] / 4096.f;
                f32AnchorY1 = pAnchor[s32LocPosi];

                s32LocPosi = ((c_loc*4 + 2) * H + h) * S_loc + w;
                //s32AnchorPosi = ((c_loc*4 + 2) * H + h) * S + w;
                f32LocW = ps32RectPtr[s32LocPosi] / 4096.f;
                f32AnchorX2 = pAnchor[s32LocPosi];

                s32LocPosi = ((c_loc*4 + 3) * H + h) * S_loc + w;
                //s32AnchorPosi = ((c_loc*4 + 3) * H + h) * S + w;
                f32LocH = ps32RectPtr[s32LocPosi] / 4096.f;
                f32AnchorY2 = pAnchor[s32LocPosi];

                f32AnchorCx = (f32AnchorX2 + f32AnchorX1)/2;
                f32AnchorCy = (f32AnchorY2 + f32AnchorY1)/2;
                f32AnchorW = f32AnchorX2 - f32AnchorX1;
                f32AnchorH = f32AnchorY2 - f32AnchorY1;

                //MP_LOG_INFO("xxxxxxxxxxxxxxxxxxxx 00  %f  %f  %f  %f \n", f32AnchorX1, f32AnchorY1, f32AnchorX2, f32AnchorY2);
                //MP_LOG_INFO("xxxxxxxxxxxxxxxxxxxx 11  %f  %f  %f  %f -- score : %f\n", f32AnchorCx, f32AnchorCy, f32AnchorW, f32AnchorH, f32Score1);

                f32BoxCx = prior_variance[0] * f32LocCx * f32AnchorW + f32AnchorCx;
                f32BoxCy = prior_variance[1] * f32LocCy * f32AnchorH + f32AnchorCy;
                f32BoxW = expf(prior_variance[2] * f32LocW) * f32AnchorW;
                f32BoxH = expf(prior_variance[3] * f32LocH) * f32AnchorH;
                //MP_LOG_INFO("xxxxxxxxxxxxxxxxxxxx 22  %f  %f  %f  %f -- score : %f\n", f32BoxCx, f32BoxCy, f32BoxW, f32BoxH, f32Score1);

                stRect.m_stRect.m_f32X1 = f32BoxCx - 0.5f * f32BoxW;
                stRect.m_stRect.m_f32Y1 = f32BoxCy - 0.5f * f32BoxH;
                stRect.m_stRect.m_f32X2 = f32BoxCx + 0.5f * f32BoxW;
                stRect.m_stRect.m_f32Y2 = f32BoxCy + 0.5f * f32BoxH;

                stRect.m_stRect.m_f32X1 = MP_MAX(0, stRect.m_stRect.m_f32X1);
                stRect.m_stRect.m_f32Y1 = MP_MAX(0, stRect.m_stRect.m_f32Y1);
                stRect.m_stRect.m_f32X2 = MP_MIN(1.f, stRect.m_stRect.m_f32X2);
                stRect.m_stRect.m_f32Y2 = MP_MIN(1.f, stRect.m_stRect.m_f32Y2);

                
                //MP_LOG_INFO("xxxxxxxxxxxxxxxxxxxx 33  %f  %f  %f  %f -- score : %f\n", 
                //    stRect.m_stRect.m_f32X1, stRect.m_stRect.m_f32Y1, 
                //    stRect.m_stRect.m_f32X2, stRect.m_stRect.m_f32Y2, f32Score1);

                stRect.m_stRect.m_f32X1 = stRect.m_stRect.m_f32X1 * 1024 + 0.5;
                stRect.m_stRect.m_f32Y1 = stRect.m_stRect.m_f32Y1 * 576 + 0.5;
                stRect.m_stRect.m_f32X2 = stRect.m_stRect.m_f32X2 * 1024 + 0.5;
                stRect.m_stRect.m_f32Y2 = stRect.m_stRect.m_f32Y2 * 576 + 0.5;
                stRect.m_f32Score = f32Score1;

                //MP_LOG_INFO("xxxxxxxxxxxxxxxxxxxx 44  %f  %f  %f  %f -- score : %f\n", 
                //    stRect.m_stRect.m_f32X1, stRect.m_stRect.m_f32Y1, 
                //    stRect.m_stRect.m_f32X2, stRect.m_stRect.m_f32Y2, f32Score1);

                if ((stRect.m_stRect.m_f32X2 - stRect.m_stRect.m_f32X1) >= 1 &&
                    (stRect.m_stRect.m_f32Y2 - stRect.m_stRect.m_f32Y1) >= 1)
                {
                    _svtRect.emplace_back(ALG_RECT_RESULT_S{ stRect });
                }
            }
        }
    };

    get_box_by_anchor(m_stvOutputSvpBlob[0], m_stvOutputSvpBlob[1], m_stvAnchor[0], m_stvRectResult);
    get_box_by_anchor(m_stvOutputSvpBlob[2], m_stvOutputSvpBlob[3], m_stvAnchor[1], m_stvRectResult);
    get_box_by_anchor(m_stvOutputSvpBlob[4], m_stvOutputSvpBlob[5], m_stvAnchor[2], m_stvRectResult);
    get_box_by_anchor(m_stvOutputSvpBlob[6], m_stvOutputSvpBlob[7], m_stvAnchor[3], m_stvRectResult);

    //MP_LOG_DBG("box number before nms:%d \n", m_stvRectResult.size());

    rect_nms(m_stvRectResult, m_f32NmsThreShold);

    return true;
}

INT32 FaceDetV1::alg_run_image(ALG_IMAGE_S* _pstImage, INT32 _s32MaxRet, void *_pstResult)
{
    //bool bFinish = false;
    INT32 s32RectSize = 0;
    UINT64 u64TimeStart = 0;
    INT32 s32CurBatch = 0;

    CHECK_NON_NULL_POINTER(_pstImage, -1);
    CHECK_NON_NULL_POINTER(_pstResult, -1);

    //SVP_IMAGE_S* pstSvpImage = (SVP_IMAGE_S*)_pstImage;
    m_input_batch = 1;
    m_stvRectResult.clear();
    
    //while (1)
    {
        s32CurBatch = 1;

        TIME_START(u64TimeStart);
        ab_image_pre(IMG_PRE_TYPE_IMAGE, _pstImage);
        TIME_END(u64TimeStart, "Alg run image pre.model:%s, batch:%d, time: ", m_cModelName, s32CurBatch);

        TIME_START(u64TimeStart);
        ab_image_proc();
        TIME_END(u64TimeStart, "Alg run image proc.model:%s, batch:%d, time: ", m_cModelName, s32CurBatch);

        TIME_START(u64TimeStart);
        if(alg_get_output_w(0) == alg_get_output_w(1))
            ab_image_post();
        else
            ab_image_post_v2();
        TIME_END(u64TimeStart, "Alg run image post.model:%s, batch:%d, time: ", m_cModelName, s32CurBatch);

        //if (bFinish)
        //{
            //break;
        //}
    }

    s32RectSize = m_stvRectResult.size() < _s32MaxRet ? m_stvRectResult.size() : _s32MaxRet;
    //optimize .
    if (s32RectSize > 0)
    {
        memcpy(_pstResult, m_stvRectResult.data(), sizeof(ALG_RECT_RESULT_S) * s32RectSize);
    }

    return s32RectSize;
}

